# This script visualizes task points colored by region assignments.
# It supports three input formats for the region encoding:
# 1) partition_result.csv     -> columns: task_id, region_id
# 2) region_binaries.json     -> keys: region_0 ... region_5, each a 0/1 list
# 3) partition_binaries.csv   -> columns: task_id, region_0 ... region_5 (0/1)
#
# It reads coordinates from the Excel file `附件一_A_data.xlsx`
# (sheet "Case1" by default), plots a single scatter with different colors
# per region (letting matplotlib pick the colors), and saves the figure to
# `region_scatter.png`.
#
# If multiple files exist, priority is: partition_result.csv > region_binaries.json > partition_binaries.csv.
#
# If none of these files exist, the script will raise a helpful error telling
# the user how to provide the encodings.
import os
import json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from pathlib import Path
from typing import List, Dict, Optional

EXCEL_PATHS = [
    "附件一_A_data.xlsx",  # uploaded path
    "附件一_A_data.xlsx"             # local alt
]
SHEET_NAME = "Case1"

# Helper: load excel
def load_xy_from_excel(paths, sheet) -> pd.DataFrame:
    last_err = None
    for p in paths:
        if os.path.exists(p):
            try:
                df = pd.read_excel(p, sheet_name=sheet)
                # Validate columns
                need = ['任务编号', 'x坐标 (km)', 'y坐标 (km)']
                for c in need:
                    if c not in df.columns:
                        raise ValueError(f"Excel缺少列: {c}")
                return df[['任务编号', 'x坐标 (km)', 'y坐标 (km)']].copy()
            except Exception as e:
                last_err = e
    if last_err:
        raise last_err
    raise FileNotFoundError("未找到 Excel：附件一_A_data.xlsx 或 当前目录下的 附件一_A_data.xlsx")

# Helper: find labels from one of the supported files
def find_labels(N: int) -> np.ndarray:
    # 1) partition_result.csv
    cand1 = ["partition_result.csv", "partition_result.csv"]
    for p in cand1:
        if os.path.exists(p):
            df = pd.read_csv(p)
            if "region_id" not in df.columns:
                raise ValueError(f"{p} 缺少列 region_id")
            labs = df["region_id"].to_numpy()
            if labs.shape[0] != N:
                raise ValueError(f"{p} 中行数 {labs.shape[0]} 与 Excel 任务数 {N} 不一致")
            return labs.astype(int)

    # 2) region_binaries.json
    cand2 = ["region_binaries.json", "region_binaries.json"]
    for p in cand2:
        if os.path.exists(p):
            with open(p, "r", encoding="utf-8") as f:
                obj = json.load(f)
            # Expect region_0 ... region_5
            keys = [k for k in obj.keys() if k.startswith("region_")]
            if not keys:
                raise ValueError(f"{p} 中未找到 region_* 键")
            K = len(keys)
            # Ensure ordering by numeric suffix
            keys = sorted(keys, key=lambda s: int(s.split("_")[-1]))
            mats = [np.array(obj[k], dtype=int) for k in keys]
            for mat in mats:
                if mat.shape[0] != N:
                    raise ValueError(f"{p}:{keys[0]} 长度与 Excel 任务数不一致")
            M = np.stack(mats, axis=1)  # (N,K)
            # Resolve to labels: take the first region with 1, else -1
            labs = np.full(N, -1, dtype=int)
            for i in range(N):
                row = M[i]
                ones = np.where(row == 1)[0]
                if ones.size > 0:
                    labs[i] = int(ones[0])
                else:
                    labs[i] = -1
            if (labs < 0).any():
                # If some tasks are unassigned, fallback to argmax to force assignment
                labs = np.argmax(M, axis=1)
            return labs.astype(int)

    # 3) partition_binaries.csv
    cand3 = ["partition_binaries.csv", "partition_binaries.csv"]
    for p in cand3:
        if os.path.exists(p):
            df = pd.read_csv(p)
            region_cols = [c for c in df.columns if c.startswith("region_")]
            if not region_cols:
                raise ValueError(f"{p} 中未找到 region_* 列")
            region_cols = sorted(region_cols, key=lambda s: int(s.split("_")[-1]))
            M = df[region_cols].to_numpy(dtype=int)  # (N,K)
            if M.shape[0] != N:
                raise ValueError(f"{p} 行数与 Excel 任务数不一致")
            labs = np.argmax(M, axis=1)
            return labs.astype(int)

    raise FileNotFoundError(
        "未找到编码文件。\n"
        "请在工作目录或 /mnt/data 下提供以下之一：\n"
        "- partition_result.csv (含列: task_id, region_id)\n"
        "- region_binaries.json (含键: region_0..region_5, 每个为0/1列表)\n"
        "- partition_binaries.csv (含列: task_id, region_0..region_5)"
    )

# Load excel & labels
df_xy = load_xy_from_excel(EXCEL_PATHS, SHEET_NAME)
N = df_xy.shape[0]
labels = find_labels(N)

# Sanity: ensure labels are 0..K-1
K = int(labels.max()) + 1
if (labels < 0).any():
    raise ValueError("存在未分配区域的任务（label<0）。请检查编码文件。")

# Build a plot
plt.figure(figsize=(8, 7))
# Plot each region as a separate series (matplotlib will auto-assign colors)
counts = []
for k in range(K):
    mask = labels == k
    counts.append(int(mask.sum()))
    plt.scatter(df_xy.loc[mask, 'x坐标 (km)'], df_xy.loc[mask, 'y坐标 (km)'], s=10, label=f"Region {k} (n={int(mask.sum())})")

plt.xlabel("x (km)")
plt.ylabel("y (km)")
plt.title("Task Locations by Region (from encodings)")
plt.legend(markerscale=1.5, fontsize=9, loc="best")
plt.grid(True, linestyle="--", linewidth=0.5, alpha=0.5)

out_path = "region_scatter.png"
plt.tight_layout()
plt.savefig(out_path, dpi=160)
plt.show()

# Also show a small summary table of counts per region
summary_df = pd.DataFrame({"region": [f"Region {i}" for i in range(K)], "count": counts})
import caas_jupyter_tools
caas_jupyter_tools.display_dataframe_to_user("Region Counts", summary_df)

out_path
